from PIL import Image
import base64
import io
from typing import Tuple, Optional, List
import hashlib
import requests
import urllib.parse
import os

class ImageProcessor:
    """图像处理器"""
    
    @staticmethod
    def image_to_base64(image_path: str) -> str:
        """将图片文件转换为base64编码"""
        try:
            with open(image_path, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
                return encoded_string
        except Exception as e:
            raise Exception(f"图片转换失败: {str(e)}")
    
    @staticmethod
    def bytes_to_base64(image_bytes: bytes) -> str:
        """将图片字节转换为base64编码"""
        try:
            encoded_string = base64.b64encode(image_bytes).decode('utf-8')
            return encoded_string
        except Exception as e:
            raise Exception(f"图片转换失败: {str(e)}")
    
    @staticmethod
    def resize_image(image_bytes: bytes, max_size: Tuple[int, int] = (1024, 1024)) -> bytes:
        """调整图片尺寸"""
        try:
            image = Image.open(io.BytesIO(image_bytes))
            
            # 如果是RGBA模式（PNG透明图片），转换为RGB模式
            if image.mode in ('RGBA', 'LA', 'P'):
                # 创建白色背景
                background = Image.new('RGB', image.size, (255, 255, 255))
                if image.mode == 'P':
                    image = image.convert('RGBA')
                background.paste(image, mask=image.split()[-1] if image.mode == 'RGBA' else None)
                image = background
            elif image.mode != 'RGB':
                image = image.convert('RGB')
            
            # 保持宽高比调整尺寸
            image.thumbnail(max_size, Image.Resampling.LANCZOS)
            
            # 转换为字节
            output = io.BytesIO()
            image.save(output, format='JPEG', quality=85)
            return output.getvalue()
            
        except Exception as e:
            raise Exception(f"图片调整失败: {str(e)}")
    
    @staticmethod
    def generate_image_hash(image_bytes: bytes) -> str:
        """生成图片哈希值"""
        return hashlib.md5(image_bytes).hexdigest()
    
    @staticmethod
    def validate_image_format(image_bytes: bytes) -> bool:
        """验证图片格式"""
        try:
            image = Image.open(io.BytesIO(image_bytes))
            return image.format.lower() in ['jpeg', 'jpg', 'png', 'bmp']
        except:
            return False
    
    @staticmethod
    def split_mj_image(image_bytes: bytes) -> List[bytes]:
        """将MJ生成的2x2拼接图分割成4张独立图片"""
        try:
            image = Image.open(io.BytesIO(image_bytes))
            
            # 如果是RGBA模式，转换为RGB模式
            if image.mode in ('RGBA', 'LA', 'P'):
                # 创建白色背景
                background = Image.new('RGB', image.size, (255, 255, 255))
                if image.mode == 'P':
                    image = image.convert('RGBA')
                background.paste(image, mask=image.split()[-1] if image.mode == 'RGBA' else None)
                image = background
            elif image.mode != 'RGB':
                image = image.convert('RGB')
            
            # 获取图片尺寸
            width, height = image.size
            
            # MJ图片通常是正方形，分割成2x2
            half_width = width // 2
            half_height = height // 2
            
            # 分割成4张图片
            images = []
            
            # 左上角
            top_left = image.crop((0, 0, half_width, half_height))
            output = io.BytesIO()
            top_left.save(output, format='JPEG', quality=95)
            images.append(output.getvalue())
            
            # 右上角
            top_right = image.crop((half_width, 0, width, half_height))
            output = io.BytesIO()
            top_right.save(output, format='JPEG', quality=95)
            images.append(output.getvalue())
            
            # 左下角
            bottom_left = image.crop((0, half_height, half_width, height))
            output = io.BytesIO()
            bottom_left.save(output, format='JPEG', quality=95)
            images.append(output.getvalue())
            
            # 右下角
            bottom_right = image.crop((half_width, half_height, width, height))
            output = io.BytesIO()
            bottom_right.save(output, format='JPEG', quality=95)
            images.append(output.getvalue())
            
            return images
            
        except Exception as e:
            raise Exception(f"图片分割失败: {str(e)}")
    
    @staticmethod
    def get_file_extension_from_url(url: str) -> Optional[str]:
        """从URL中提取文件扩展名"""
        try:
            import urllib.parse
            import os
            
            # 解析URL
            parsed_url = urllib.parse.urlparse(url)
            path = parsed_url.path
            
            # 提取文件扩展名
            _, ext = os.path.splitext(path)
            
            if ext:
                return ext.lower()
            
            # 如果URL中没有扩展名，尝试从Content-Type判断
            try:
                response = requests.head(url, timeout=10)
                content_type = response.headers.get('content-type', '').lower()
                
                if 'jpeg' in content_type or 'jpg' in content_type:
                    return '.jpg'
                elif 'png' in content_type:
                    return '.png'
                elif 'webp' in content_type:
                    return '.webp'
                elif 'gif' in content_type:
                    return '.gif'
                else:
                    return '.jpg'  # 默认扩展名
            except:
                return '.jpg'  # 默认扩展名
                
        except Exception as e:
            print(f"提取文件扩展名失败: {str(e)}")
            return '.jpg'  # 默认扩展名
    
    @staticmethod
    def download_image_from_url(url: str) -> bytes:
        """从URL下载图片"""
        try:
            response = requests.get(url, timeout=30)
            response.raise_for_status()
            return response.content
        except Exception as e:
            raise Exception(f"下载图片失败: {str(e)}")

    @staticmethod
    def to_jpeg(image_bytes: bytes, quality: int = 85) -> bytes:
        """将任意格式的图片字节转换为JPEG，保持原始尺寸"""
        try:
            image = Image.open(io.BytesIO(image_bytes))
            # 处理带透明通道/调色板的图片，转换为RGB白底
            if image.mode in ('RGBA', 'LA', 'P'):
                background = Image.new('RGB', image.size, (255, 255, 255))
                if image.mode == 'P':
                    image = image.convert('RGBA')
                background.paste(image, mask=image.split()[-1] if image.mode == 'RGBA' else None)
                image = background
            elif image.mode != 'RGB':
                image = image.convert('RGB')

            output = io.BytesIO()
            image.save(output, format='JPEG', quality=quality)
            return output.getvalue()
        except Exception as e:
            raise Exception(f"图片转换为JPEG失败: {str(e)}")